// clust.c
// -------
// CS181 PS3 - Clustering and Parameter Estimation
// Eric Huang
// Bohao (Dan) Pan

#include <stdio.h>
#include <string>
#include <stdlib.h>
#include <iostream>
#include <time.h>
#include "clust_util.h"
#include <math.h>
#include <vector>
#include <sstream>
#include <cmath>
using std::stringstream;

//define DATAFILE "adults-small.dat"

using namespace std;

void Cluster_Kmeans(DatasetDescription DD, inst *dataset, int num_examples, int num_clusters);
double dist_sq(inst a, inst b);

void Cluster_HAC(DatasetDescription DD, inst *dataset, int num_examples, int num_clusters, double (*dist_func)(vector<inst> a, vector<inst> b));
double HAC_dist_min(vector<inst> a, vector<inst> b);
double HAC_dist_max(vector<inst> a, vector<inst> b);
double HAC_dist_mean(vector<inst> a, vector<inst> b);
double HAC_dist_centroid(vector<inst> a, vector<inst> b);
void HAC_print_centroid(vector<inst> a);
inst* HAC_get_centroid(vector<inst> a);
void HAC_print_all(vector<inst> a);

void Cluster_Autoclass(DatasetDescription DD, inst *dataset, int num_examples, int num_clusters);


// main
// ----
// The main program loop
// You should modify this function to run your experiments

int main(int argc, char **argv)
{
  
  int num_clusters;	// Number of clusters to use 
  inst *dataset;
  FILE *datafile;	// The datafile being used
  int num_examples_to_use; // Number of instances
  string clust_type;
  
  srand48(time(NULL));
  DatasetDescription DD; // The DatasetDescription for this dataset
  
  // Parse the command line
  if (argc == 4) {  // custom args with flag for clustering algorithm
    int valid = 1;
    num_clusters = atoi(argv[1]);
    if (num_clusters == 0) {
      valid = 0;
    }
    num_examples_to_use = atoi(argv[2]);
    if (num_examples_to_use == 0) {
      valid = 0;
    }
    clust_type = argv[3];
    if (clust_type != "kmeans" && 
        clust_type != "hac-min" && 
        clust_type != "hac-max" && 
        clust_type != "hac-mean" &&
        clust_type != "hac-centroid" && 
        clust_type != "autoclass" ) {
      valid = 0;
    }
    if (valid == 0) {
      cout << "usage: clust num_clusters num_examples [kmeans|hac-(min|max|mean|centroid)|autoclass]" << endl;
      exit(1);
    }
  }
  else {  // original
    if (ValidateInput(argc, argv, &num_clusters, &num_examples_to_use) < 0) {
      cout << "usage: clust num_clusters num_examples [kmeans|hac-(min|max|mean|centroid)|autoclass]" << endl;
      exit(1);
    }
  } 
  
  // defaults
  string DATAFILE = "adults.dat";
  DD.numattribs = 48;
  
  // Set DD to describe this particular dataset
  if (clust_type == "kmeans" || clust_type == "autoclass") {
    DD.numattribs = 48; // Num Attributes = 48 for adults dataset
    DATAFILE = "adults.dat";
  }
  else { //if (clust_type.find("hac-")==0) {
    DD.numattribs = 3; // Num Attributes = 3 for adults-small dataset
    DATAFILE = "adults-small.dat";
  }
  
  dataset = (inst *)malloc(num_examples_to_use*sizeof(inst));
  // Read in the data file
  datafile = fopen(DATAFILE.c_str(), "r");
  
  if(datafile == NULL) {
    cout << "Unable to open data file" << endl;
	exit(1);
  }

  if(ParseInput(DD, dataset, num_examples_to_use, datafile) < 0) {
    cout << "Error reading data file" << endl;
	exit(1);
  }
  
  fclose(datafile);
  
  if (clust_type == "kmeans") {
    Cluster_Kmeans(DD, dataset, num_examples_to_use, num_clusters);
  }
  else if (clust_type == "hac-min") {
    Cluster_HAC(DD, dataset, num_examples_to_use, num_clusters, &HAC_dist_min);
  }
  else if (clust_type == "hac-max") {
    Cluster_HAC(DD, dataset, num_examples_to_use, num_clusters, &HAC_dist_max);
  }
  else if (clust_type == "hac-mean") {
    Cluster_HAC(DD, dataset, num_examples_to_use, num_clusters, &HAC_dist_mean);
  }
  else if (clust_type == "hac-centroid") {
    Cluster_HAC(DD, dataset, num_examples_to_use, num_clusters, &HAC_dist_centroid);
  }
  else if (clust_type == "autoclass") {
    Cluster_Autoclass(DD, dataset, num_examples_to_use, num_clusters);
  }
  
  return 0;
}

void Cluster_Kmeans(DatasetDescription DD, inst *dataset, int num_examples, int num_clusters) {
  inst *mu = (inst *)malloc(num_clusters*sizeof(inst));  
  int *r = (int *)malloc(num_examples*sizeof(int));
  int *rprev = (int *)malloc(num_examples*sizeof(int));
  double error = INT_MAX;
  
  // Set mu_k to random vector in dataset
  srand( 1 ); // constant seed for consistency in testing
  for (int k=0; k<num_clusters; k++) {
    mu[k] = dataset[rand()%num_examples];
    //mu[k] = dataset[k];
  }

  // Repeat until convergence
  for (int converged=0; converged==0; ) {
    // E-step
    // For each i, set r[i] = arg min(k', dist_sq(dataset[i],mu[k]))
    for (int i=0; i<num_examples; i++) {
      double min_dist = INT_MAX;    
      for (int k=0; k<num_clusters; k++) {
        double dist = dist_sq(dataset[i], mu[k]);
        if (dist < min_dist) {
          r[i] = k;
          min_dist = dist;
        }
      }
    }
    
    // M-step
    // For each k, set mu[k] = mean(dataset[i] s.t. r[i]==k)
    for (int k=0; k<num_clusters; k++) {
      double sum[48] = {0};
      int n = 0;
      for (int i=0; i<num_examples; i++) {
        if (r[i] == k) {
          for (int z=0; z<48; z++) 
            sum[z] += dataset[i].attribs[z];            
          n++;
        }
      }
      for (int z=0; z<48; z++)
        mu[k].attribs[z] = sum[z]/n;      
    }
    
    // Error calculation
    // Err = Sum(dist_sq(dataset[i], mu[r[i]]))
    error = 0;
    for (int i=0; i<num_examples; i++) {
      error += dist_sq(dataset[i], mu[r[i]]) / num_examples;
    }
    //cout << error << endl;
  
    // Check for convergence
    converged = 1;
    for (int i=0; i<num_examples; i++) {
      if (rprev[i] != r[i]) {
        converged = 0;
      }
      rprev[i] = r[i];
    }
  }
  
  // Print mean squared error
  cout << "Mean Squared Error: " << error << endl;
  
  // Print means of clusters
  //PrintOutput(DD, mu, num_clusters);
  int *clusterSize = (int *)malloc(num_clusters*sizeof(int));
  for (int k=0; k<num_clusters; k++) clusterSize[k] = 0;  
  for (int i=0; i<num_examples; i++)
    clusterSize[r[i]]++;  
  for (int k=0; k<num_clusters; k++) {
    cout << "Cluster " << k << ": (" << clusterSize[k] << ") " << endl;
    PrintOutput(DD, &mu[k], 1);
  }
}

double dist_sq(inst a, inst b) {
  double dist_sq = 0;
  for(int i=0; i<48; i++) {
    dist_sq += (a.attribs[i]-b.attribs[i])*(a.attribs[i]-b.attribs[i]);
  }
  return dist_sq;
}

void Cluster_HAC(DatasetDescription DD, inst *dataset, int num_examples, int num_clusters, double (*dist_func)(vector<inst> a, vector<inst> b)) {
  
  // B = {dataset[i] for all i}
  vector< vector<inst> > B;
  for (int i=0; i<num_examples; i++) {
    vector<inst> c;
    c.push_back(dataset[i]);
    B.push_back(c);
  }
  
  // Repeat until |B| <= K
  while ((int)B.size() > num_clusters) {
    
    // Find two closest clusters (min_dist_ci, min_dist_cj) in B
    double min_dist = INT_MAX;
    vector< vector<inst> >::iterator min_dist_ci, min_dist_cj;
    for (vector< vector<inst> >::iterator iti = B.begin(); iti != B.end()-1; ++iti) {
      for (vector< vector<inst> >::iterator itj = iti+1; itj != B.end(); ++itj) {
        double dist = (*dist_func)(*iti, *itj);
        if (dist < min_dist) {
          min_dist = dist;
          min_dist_ci = iti;
          min_dist_cj = itj;
        }
      }
    }
    
    // Insert (min_dist_ci U min_dist_cj) into B
    vector<inst> cunion;
    cunion.insert(cunion.end(), min_dist_ci->begin(), min_dist_ci->end());
    cunion.insert(cunion.end(), min_dist_cj->begin(), min_dist_cj->end());
    B.push_back(cunion);
    
    // Remove two closest clusters (min_dist_ci, min_dist_cj) from B
    // Careful! vector changes depending on where we remove first
    // We are guaranteed that cj > ci, so this is ok
    B.erase(min_dist_cj);
    B.erase(min_dist_ci);
  }
  
  // Print summary of B  
  cout << "cluster\tv1\tv2\tv3" << endl;
  for (vector< vector<inst> >::iterator it = B.begin(); it != B.end(); ++it) {
    //cout << "Cluster " << it - B.begin() << " @ " ;
    //HAC_print_centroid(*it);
    //HAC_print_all(*it);
    
    // Prints out (cluster, points)
    for (vector<inst>::iterator itv = (*it).begin(); itv != (*it).end(); ++itv) {
      cout << it-B.begin() << " ";
      for (int i=0; i<3; i++) {
        cout << itv->attribs[i] << " ";
      }
      cout << endl;
    }
  }
}

double HAC_dist_min(vector<inst> a, vector<inst> b) {
  double mindist = INT_MAX;
  for (vector<inst>::iterator ita = a.begin(); ita != a.end(); ++ita) {
    for (vector<inst>::iterator itb = b.begin(); itb != b.end(); ++itb) {
      double currdist = dist_sq(*ita, *itb);
      if (currdist < mindist) {
        mindist = currdist;
      }
    }
  }
  return mindist;
}

double HAC_dist_max(vector<inst> a, vector<inst> b) {
  double maxdist = 0;
  for (vector<inst>::iterator ita = a.begin(); ita != a.end(); ++ita) {
    for (vector<inst>::iterator itb = b.begin(); itb != b.end(); ++itb) {
      double currdist = dist_sq(*ita, *itb);
      if (currdist > maxdist) {
        maxdist = currdist;
      }
    }
  }
  return maxdist;
}

double HAC_dist_mean(vector<inst> a, vector<inst> b) {
  double dist = 0;
  for (vector<inst>::iterator ita = a.begin(); ita != a.end(); ++ita) {
    for (vector<inst>::iterator itb = b.begin(); itb != b.end(); ++itb) {
      dist += sqrt(dist_sq(*ita, *itb))/(a.size()*b.size());
    }
  }
  return dist;
}

double HAC_dist_centroid(vector<inst> a, vector<inst> b) {
  inst *centroida = HAC_get_centroid(a);
  inst *centroidb = HAC_get_centroid(b);
  return dist_sq(*centroida, *centroidb);
}

void HAC_print_centroid(vector<inst> a) {
  inst* centroid = HAC_get_centroid(a);
  cout << "(" << a.size() << ") " ;  
  for (int i=0; i<3; i++) {
    cout << centroid->attribs[i] << ",";
  }
  cout << endl;
}

inst* HAC_get_centroid(vector<inst> a) {
  inst *centroid = (inst *)malloc(sizeof(inst));
  for(int i=0; i<48; i++) {
    centroid->attribs[i] = 0;
  }
  for (vector<inst>::iterator ita = a.begin(); ita != a.end(); ++ita) {
    for (int i=0; i<3; i++) {
      centroid->attribs[i] += ita->attribs[i]/a.size();
    }
  }
  return centroid;
}

void HAC_print_all(vector<inst> a) {
  for (vector<inst>::iterator ita = a.begin(); ita != a.end(); ++ita) {
    for (int i=0; i<3; i++) {
      cout << ita->attribs[i] << ",";
    }
    cout << ";";
  }
  cout << endl;
}

void Cluster_Autoclass(DatasetDescription DD, inst *dataset, int num_examples, int num_clusters) {
	int num_buckets = 5;
	
	double *theta_c = (double *)malloc((num_clusters)*sizeof(double));
	double *theta_cprev = (double *)malloc((num_clusters)*sizeof(double));
	
	// theta[j][b][k] is Pr(X_j is in bucket b | C = k)
	double ***theta = (double ***)malloc((DD.numattribs)*sizeof(double **));
	double ***thetaprev = (double ***)malloc((DD.numattribs)*sizeof(double **));
	for (int j = 0; j < DD.numattribs; j++) {
		theta[j] = (double **)malloc(num_buckets*sizeof(double *));
		thetaprev[j] = (double **)malloc(num_buckets*sizeof(double *));
		for (int b = 0; b < num_buckets; b++) {
			theta[j][b] = (double *)malloc((num_clusters)*sizeof(double));
			thetaprev[j][b] = (double *)malloc((num_clusters)*sizeof(double));
		}
	}
	
	// initialize theta
  srand( 1000 ); // constant seed for consistency in testing

	double sum = 0;
	for (int i = 0; i < num_clusters - 1; i++) {
		theta_c[i] = theta_cprev[i] = 1.0/num_clusters;
		sum += theta_c[i];
	}
	theta_c[num_clusters-1] = theta_cprev[num_clusters-1] = 1.0 - sum;

	for (int j = 0; j < DD.numattribs; j++) {
		for (int b = 0; b < num_buckets; b++) {
			for (int k = 0; k < num_clusters; k++) {
				theta[j][b][k] = thetaprev[j][b][k] = (double) (rand()%1000)/1000;
			}
		}
	}
	
	// N[k] = number of instances classified as k
	double *N = (double *)malloc((num_clusters)*sizeof(double));
	
	// NN[j][b][k] = number of instances i such that t_i = k and x_ij is in bucket b
	double ***NN = (double ***)malloc(num_examples*sizeof(double **));
	for (int j = 0; j < DD.numattribs; j++) {
		NN[j] = (double **)malloc(num_buckets*sizeof(double *));
		for (int b = 0; b < num_buckets; b++) {
			NN[j][b] = (double *)malloc(num_clusters*sizeof(double));
		}
	}
	
	// p's to be used in expectation step
	// p[k] is Pr(x_i | C_i = k)*Pr(C_i = k)
	double *p = (double *)malloc(num_clusters*sizeof(double));
	
	// pc[i][k] is Pr(C_i = k | x_i)
	double **pc = (double **)malloc(num_examples*sizeof(double *));
	for (int i = 0; i < num_examples; i++) {
		pc[i] = (double *)malloc(num_clusters*sizeof(double));
	}
	
	// p_x[i] is Pr(x_i)
	double *p_x = (double *)malloc(num_examples*sizeof(double));
  
	/////////////////////////////////////////
  // Repeat until convergence
	int iter = 0;
  for (int converged=0; converged==0; iter++) {
		///////////////////////
    // E-step
    
		// reset expectations
		for (int i = 0; i < num_clusters; i++) {
			N[i] = 0.0;
		}
		for (int j = 0; j < DD.numattribs; j++) {
			for (int b = 0; b < num_buckets; b++) {
				for (int k = 0; k < num_clusters; k++) {
					NN[j][b][k] = 0.0;
				}
			}
		}
    // For each instance i
    for (int i = 0; i < num_examples; i++) {
			// for each cluster k, calculate p_k
			p_x[i] = 0.0;
			for (int k = 0; k < num_clusters; k++) {
				p[k] = theta_c[k];
				// for each attribute j
				for (int j = 0; j < DD.numattribs; j++) {
					int b = (int) (dataset[i].attribs[j]*num_buckets);
					b = (b == num_buckets)? num_buckets-1 : b; // edge case
					p[k] *= theta[j][b][k];
				}
				p_x[i] += p[k];
    	}
      
			// update expectations
  		for (int k = 0; k < num_clusters; k++) {
				pc[i][k] = p[k]/p_x[i];
  			N[k] += pc[i][k];
	
        for (int j = 0; j < DD.numattribs; j++) {
					int b = (int) (dataset[i].attribs[j]*num_buckets);
					b = (b == num_buckets)? num_buckets-1 : b; //edge case
					NN[j][b][k] += pc[i][k];
				}
			}
		}
    
		///////////////////////////
    // M-step
		for (int k = 0; k < num_clusters; k++) {
			theta_c[k] = N[k]/num_examples;
			
			for (int j = 0; j < DD.numattribs; j++) {
				for (int b = 0; b < num_buckets; b++) {
					theta[j][b][k] = NN[j][b][k]/N[k];
				}
			}
		}
		
		////////////////////////////////////
		// compute log likelihood of data
		// P(D|theta) = P(x_1|theta)...P(x_n|theta)
		double lll = 0.0;
		for (int i = 0; i < num_examples; i++) {
			lll += log(p_x[i]);
		}
		cout << iter << " " << "Log Likelihood: " <<lll <<endl;

		////////////////////////////
    // Check for convergence
    converged = 1;
		for (int k = 0; k < num_clusters; k++) {
			if (abs(theta_cprev[k] - theta_c[k]) > 1.0e-8)
				converged = 0;
			theta_cprev[k] = theta_c[k];
			
			for (int j = 0; j < DD.numattribs; j++) {
				for (int b = 0; b < num_buckets; b++) {
					if (abs(thetaprev[j][b][k] - theta[j][b][k]) > 1.0e-8)
						converged = 0;
					thetaprev[j][b][k] = theta[j][b][k];
				}
			}
		}
  }
  
  cout << "Number of iterations (to get to 1.0e-8 precision for all theta): " << iter <<endl;
	
	for (int k = 0; k < num_clusters; k++) {
		for (int j = 0; j < DD.numattribs; j++) {
			for (int b = 0; b < num_buckets; b++) {
				cout << theta[j][b][k] << " ";
			}
			cout <<endl;
		}
	}
	////////////////////////////////
	// get means of each cluster
  inst *mu = (inst *)malloc(num_clusters*sizeof(inst));
  int *r = (int *)malloc(num_examples*sizeof(int));
  
	// for each cluster
	for (int k = 0; k < num_clusters; k++) {
		double sum[48] = {0};
		int n = 0;
		for (int i = 0; i < num_examples; i++) {
			// classify to the largest probability
			double max = INT_MIN;
			int cluster = -1;
			for (int kk = 0; kk < num_clusters; kk++) {
				if (pc[i][kk] > max) {
					max = pc[i][kk];
					cluster = kk;
          r[i] = kk;
				}
			}
      
			if (cluster == k) {
				n++;
				for (int z = 0; z < DD.numattribs; z++)
					sum[z] += dataset[i].attribs[z];
			}
		}	
		for (int z = 0; z < DD.numattribs; z++)
	  	mu[k].attribs[z] = sum[z]/n;
	}
  
  // mean squared error calculation for sanity check
  double error = 0.0;
  int *clusterSize = (int *)malloc(num_clusters*sizeof(int));
  for (int k=0; k<num_clusters; k++) clusterSize[k] = 0;  
  for (int i=0; i<num_examples; i++) {
    error += dist_sq(mu[r[i]], dataset[i]) / num_examples;
    clusterSize[r[i]]++;
  }
  cout << "Mean squared error: " << error << endl;
  
  for (int k=0; k<num_clusters; k++) {
    cout << "Cluster " << k << ": (" << clusterSize[k] << ") " << endl;
    PrintOutput(DD, &mu[k], 1);
  }
  
	/*
  // Print means of clusters
	PrintOutput(DD, mu, num_clusters);
  */
}
